// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/caching/ssd_file.h>

namespace kumo::pollux::cache {
#define POLLUX_SSD_CACHE_LOG_PREFIX "[SSDCA] "
#define POLLUX_SSD_CACHE_LOG(severity) \
  KLOG(severity) << POLLUX_SSD_CACHE_LOG_PREFIX

    namespace test {
        class SsdCacheTestHelper;
    }

    class SsdCache {
    public:
        struct Config {
            Config() = default;

            Config(
                const std::string &_filePrefix,
                uint64_t _maxBytes,
                int32_t _numShards,
                melon::Executor *_executor,
                uint64_t _checkpointIntervalBytes = 0,
                bool _disableFileCow = false,
                bool _checksumEnabled = false,
                bool _checksumReadVerificationEnabled = false)
                : filePrefix(_filePrefix),
                  maxBytes(_maxBytes),
                  numShards(_numShards),
                  checkpointIntervalBytes(_checkpointIntervalBytes),
                  disableFileCow(_disableFileCow),
                  checksumEnabled(_checksumEnabled),
                  checksumReadVerificationEnabled(_checksumReadVerificationEnabled),
                  executor(_executor) {
            };

            std::string filePrefix;
            uint64_t maxBytes;
            int32_t numShards;

            /// Checkpoint after every 'checkpointIntervalBytes'/'numShards' written
            /// into each file. 0 means no checkpointing.
            uint64_t checkpointIntervalBytes;

            /// True if copy on write should be disabled.
            bool disableFileCow;

            /// If true, checksum write to SSD is enabled.
            bool checksumEnabled;

            /// If true, checksum read verification from SSD is enabled.
            bool checksumReadVerificationEnabled;

            /// Executor for async fsync in checkpoint.
            melon::Executor *executor;

            std::string toString() const {
                return fmt::format(
                    "{} shards, capacity {}, checkpoint size {}, file cow {}, checksum {}, read verification {}",
                    numShards,
                    succinctBytes(maxBytes),
                    succinctBytes(checkpointIntervalBytes),
                    (disableFileCow ? "DISABLED" : "ENABLED"),
                    (checksumEnabled ? "ENABLED" : "DISABLED"),
                    (checksumReadVerificationEnabled ? "ENABLED" : "DISABLED"));
            }
        };

        /// Constructs a cache with backing files at path 'filePrefix'.<ordinal>.
        /// <ordinal> ranges from 0 to 'numShards' - 1.
        /// 'maxBytes' is the total capacity of the cache. This is rounded up to the
        /// next multiple of kRegionSize * 'numShards'. This means that all the shards
        /// have an equal number of regions. For 2 shards and 200MB size, the size
        /// rounds up to 256M with 2 shards each of 128M (2 regions).
        /// If 'checkpointIntervalBytes' is non-zero, the cache makes a durable
        /// checkpointed state that survives restart after each
        /// 'checkpointIntervalBytes' written.
        /// If 'disableFileCow' is true, the cache disables the file COW (copy on
        /// write) feature if the underlying filesystem (such as brtfs) supports it.
        /// This prevents the actual cache space usage on disk from exceeding the
        /// 'maxBytes' limit and stop working.
        SsdCache(const Config &config);

        /// Returns the shard corresponding to 'fileId'. 'fileId' is a file id from
        /// e.g. FileCacheKey.
        SsdFile &file(uint64_t fileId);

        /// Returns the maximum capacity, rounded up from the capacity passed to the
        /// constructor.
        uint64_t maxBytes() const {
            return files_[0]->maxRegions() * SsdFile::kRegionSize * files_.size();
        }

        /// Returns true if no write is in progress. Atomically sets a write to be in
        /// progress. write() must be called after this. The writing state is reset
        /// asynchronously after writing to SSD finishes.
        bool startWrite();

        bool writeInProgress() const {
            return writesInProgress_ != 0;
        }

        /// Stores the entries of 'pins' into the corresponding files. Sets the file
        /// for the successfully stored entries. May evict existing entries from
        /// unpinned regions.
        ///
        /// NOTE: startWrite() must have been called first and it must have returned
        /// true.
        void write(std::vector<CachePin> pins);

        /// Invoked to write checkpoints to all ssd files. This is used by Prestissimo
        /// worker operation.
        ///
        /// NOTE: startWrite() must have been called first and it must have returned
        /// true.
        void checkpoint();

        /// Removes cached entries from all SsdFiles for files in the fileNum set
        /// 'filesToRemove'. If successful, return true, and 'filesRetained' contains
        /// entries that should not be removed, ex., from pinned regions. Otherwise,
        /// return false and 'filesRetained' could be ignored.
        bool removeFileEntries(
            const melon::F14FastSet<uint64_t> &filesToRemove,
            melon::F14FastSet<uint64_t> &filesRetained);

        /// Returns stats aggregated from all shards.
        SsdCacheStats stats() const;

        FileGroupStats &groupStats() const {
            return *groupStats_;
        }

        /// Stops writing to the cache files and waits for pending writes to finish.
        /// If checkpointing is on, makes a checkpoint.
        void shutdown();

        std::string toString() const;

        const std::string &filePrefix() const {
            return filePrefix_;
        }

        /// Drops all entries. Outstanding pins become invalid but reading them will
        /// mostly succeed since the files will not be rewritten until new content is
        /// stored.
        ///
        /// NOTE: it is used by test and Prestissimo worker operation.
        void clear();

        /// Waits until the pending ssd cache writes or checkpoints to finish. Used by
        /// test and Prestissimo worker operation.
        void waitForWriteToFinish();

    private:
        void checkNotShutdownLocked() {
            POLLUX_CHECK(
                !shutdown_, "Unexpected write after SSD cache has been shutdown");
        }

        const std::string filePrefix_;
        const int32_t numShards_;
        // Stats for selecting entries to save from AsyncDataCache.
        const std::unique_ptr<FileGroupStats> groupStats_;
        melon::Executor *const executor_;
        mutable std::mutex mutex_;

        std::vector<std::unique_ptr<SsdFile> > files_;

        // Count of shards with unfinished writes.
        std::atomic_int32_t writesInProgress_{0};
        bool shutdown_{false};

        friend class test::SsdCacheTestHelper;
    };
} // namespace kumo::pollux::cache
